import torch
from transformers import AutoModelForCausalLM
from transformers import AutoTokenizer
import util
from datetime import datetime
from accelerate import Accelerator

accelerator = Accelerator()

torch.manual_seed(1234)
model_name = "hf-models/Qwen-VL-Chat-Int4"
# Initialize the tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
# Pass the default decoding hyperparameters of Qwen2-7B-Instruct
# max_tokens is for the maximum length for generation.
# Input the model name or path. Can be GPTQ or AWQ models.
model = AutoModelForCausalLM.from_pretrained(
    model_name, 
    torch_dtype=torch.bfloat16,
    device_map="cuda", 
    trust_remote_code=True).eval()

def recognize(image_url, target_text):
    # Prepare your prompts
    example_json = {"产品名称": "新型冠状病毒(2019-nCoV)抗原检测试剂盒（胶体金法）",  "生产商": "北京乐普诊断科技股份有限公司" }
    prompt = f'''
        ## 角色：
        你是一个高效的智能助手，可以精准、广泛地根据提取目标从图片中识别相应信息。
        ## 任务：
        你现在的任务是从图片中提取相关信息，提取目标包括：{target_text}，各提取目标使用|分隔，请自行拆分，可寻找相似目标，但禁止随意识别不相关信息。
        ## 步骤：
        1. 将提取目标 {target_text} 根据“|”符号进行分隔，并分组管理；
        2. 针对每组的目标，分析相近短语，方便后续识别时更加泛化；
        3. 针对每组的目标，按照后续“要求”进行识别，要考虑“步骤2”中的相近短语；
        4. 按照要求进行结果返回。
        参考示例：
        当提取目标为“产品名称|生产商”时，识别结果为：
        ```json{example_json}```
        ## 要求：
        1. 返回结果使用json格式，每条数据包含一个key-value对，key值为我指定的关键信息，value值为所抽取的结果。
        2. 如果认为OCR识别结果中没有关键信息key，则将value赋值为“未找到相关信息”即可。如果图片模糊，请不要随意猜测内容，将value赋值为“无法识别”即可。 
        3. 直接输出json格式的结果，严禁输出其他内容，包括但不限于：提示语、注释、符号等。
        '''
    util.download_and_show_image(image_url)
    print("image_url", str(datetime.now()) + image_url)
    query = tokenizer.from_list_format([
        {'image': util.image_path},
        {'text': prompt},
    ])
    response, history = model.chat(tokenizer, query=query, history=None)
    print(str(datetime.now()) + response)
    return response